#include <iostream>
#include <fstream>
#include <string>
#include <cstdlib>
#include <algorithm>
#include <iterator>
#include <cmath>
#include "naiev_bayes.h"

using namespace std;

//because sqrt and exp function in *.h file cause problem.
//Put the function here.
float
calc (float val, float mean, float variance2)
{
    float sqrt_2pi = 2.5066283;
    float variance = sqrt ((float) variance2);
    float exp_val = -((val - mean) * (val - mean)) / (2 * variance2);
    if (variance == 0)
    {
        cout << "error:overfloat." << endl;
        exit (1);
    }
    return (1 / (sqrt_2pi * variance)) * exp ((float) exp_val);
}

void print_version()
{
    string VERSION="jdm(Junix Data Miner) 0.01\n\
 Copyright Junix(unix_jun@yahoo.com.cn)\n\
 This is free software.\n ";
    cout<<VERSION;
}

void print_help()
{
    string USAGE="Usage:   jdm [OPTION [ARGUMENT]] ...\n\
Example: jdm -c test.date\n\n\
 -V  --version            Show the version infomation.\n\
 -h  --help               Show help infomation.\n\
 -l  --learn lf           Use lf as learning data.\n\
 -a  --attr af            Use af as the attribution infomation\n\
 -s  --attr_select asf    Use asf as selected attributions configure file.\n\
 -t  --test tf            Use tf as test data.\n\
 -c  --classify cf        Use cf as classify data.\n\n\
Report bugs to <unix_jun@yahoo.com.cn>\n ";
    cout<<USAGE;
}

int main (int argc, char **argv)
{
    //default case
    string select_attrs_conf = "select_attrs.conf";
    string attr_infos_conf = "attr_infos.conf";
    string training_data = "training.data";
    string test_data = "training.data";
    string classify_data = "test.data";

    bool test = false;
    bool classify = false;

    if (argc == 1)
    {
        print_help();
        return 0;
    }

    int arg_idx = 1;
    while (arg_idx < argc)
    {
         char* arg = argv[arg_idx];
        if (strcmp(arg, "-V") == 0 || strcmp(arg ,"--version") == 0)
        {
            print_version();
            arg_idx++;
            return 0;
        }
        else  if ( strcmp(arg,"-h") == 0 || strcmp(arg ,"--help") == 0)
        {
            print_help();
            arg_idx++;
            return 0;
        }
        else if (strcmp(arg ,"--attr_select") == 0 || strcmp(arg , "-s") == 0)
        {
            arg_idx++;
            if (arg_idx < argc)
            {
                cout<<"attribution selecting configure file is needed."<<std::endl;
                return 1;
            }
            else
            {
                string sa(argv[arg_idx]);
                select_attrs_conf = sa;
                arg_idx ++;
            }
        }
        else if ( strcmp(arg, "--attr") == 0 || strcmp(arg ,"-a") == 0 )
        {
            arg_idx++;
            if (arg_idx >= argc)
            {
                cout<<"attribution infomation file is needed."<<std::endl;
                return 1;
            }
            else
            {
                string info(argv[arg_idx]);
                attr_infos_conf = info;
                arg_idx ++;
            }

        }
        else if ( strcmp(arg ,"--learn") == 0 ||strcmp(arg, "-l") == 0 )
        {
            arg_idx++;
            if (arg_idx >= argc)
            {
                cout<<"learning data file is needed."<<std::endl;
                return 1;
            }
            else
            {
                string info(argv[arg_idx]);
                training_data = info;
                arg_idx ++;
            }

        }
        else if ( strcmp(arg, "--test") ==0 || strcmp(arg, "-t") == 0 )
        {
	    test = true;
            arg_idx++;
            if (arg_idx >= argc)
            {
                cout<<"test data file is needed."<<std::endl;
                return 1;
            }
            else
            {
                string info(argv[arg_idx]);
                test_data = info;
                arg_idx ++;
            }

        }
        else if ( strcmp( arg , "--classify") ==0 || strcmp(arg, "-c") == 0 )
        {
	    classify = true;
            arg_idx++;
            if (arg_idx >= argc)
            {
                cout<<"classify data file is needed."<<std::endl;
                return 1;
            }
            else
            {
                string info(argv[arg_idx]);
                classify_data = info;
                arg_idx ++;
            }

        }
        else
        {
            print_help();
	    return 0;
        }
    }


    select_attrs sel (select_attrs_conf);
    sel.print();
    vector < attribute_info > *infos =
        attribute_info::get_all_attr_info (attr_infos_conf);
    cout<<endl<<endl;
    copy(infos->begin(),infos->end(),ostream_iterator<attribute_info>(cout));

    instances *learn_set = new instances (training_data, sel.bitmap (), infos);
    learn_set->ignore_missing ();

    naive_bayes bayes;
    bayes.cont_cond_prob = calc;
    bayes.set_training_set (learn_set);
    bayes.learn ();
   cout<<bayes<<endl;
   
    if(classify)
    { 
       instances *classify_set = new instances (classify_data, sel.bitmap (), infos);
       classify_set->data_clean ();

       bayes.set_classify_set (classify_set);
       bayes.classify ();
       //ugly code.Special for NJU DM Class format
       bayes.nju_format();
       ofstream of("output.txt");
       bayes.print_classified_set (of);
    }

    if(test)
    {
       instances *test_set = new instances (test_data, sel.bitmap (), infos);
       test_set->data_clean ();

       bayes.set_test_set (test_set);
       bayes.test();
    }

    return 0;
}
